import sys
import rasterio
import numpy as np
from PyQt5.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QMessageBox
from qt_05RasterView_ui import Ui_RasterView
from PyQt5.QtWidgets import QFileDialog
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import matplotlib.pyplot as plt


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.ui = Ui_RasterView()
        self.ui.setupUi(self)
        # 连接按钮的点击信号到对应的槽函数（自定义的打开文件函数）
        self.ui.open_raster_button.clicked.connect(self.open_raster_data)
        # 连接下拉列表的当前索引改变信号到对应的槽函数（用于实时更新图像显示）
        self.ui.band_combo_box_1.currentIndexChanged.connect(self.update_image)
        self.ui.band_combo_box_2.currentIndexChanged.connect(self.update_image)
        self.ui.band_combo_box_3.currentIndexChanged.connect(self.update_image)

    def open_raster_data(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择遥感栅格数据文件", "", "GeoTIFF (*.tif);;ENVI (*.dat)")
        if file_path:
            try:
                self.raster_data = rasterio.open(file_path)
                self.band_count = self.raster_data.count
                # 动态添加波段序号到下拉列表
                self.populate_band_combo_boxes()
                self.ui.info_label.setText(f'成功打开遥感栅格数据，包含 {self.band_count} 个波段。')
                # 显示数据路径到lineedit控件
                self.ui.path_line_edit.setText(file_path)
                # 初始显示图像（使用默认的波段组合，可根据实际调整）
                self.update_image()
            except Exception as e:
                self.ui.info_label.setText(f'打开文件出现错误: {str(e)}')
        else:
            self.ui.info_label.setText('未选择任何文件。')

    def populate_band_combo_boxes(self):
        """
        根据栅格数据的波段数量，动态添加波段序号到三个下拉列表
        """
        band_indices = list(range(1, self.band_count + 1))
        self.ui.band_combo_box_1.clear()
        self.ui.band_combo_box_2.clear()
        self.ui.band_combo_box_3.clear()
        for index in band_indices:
            self.ui.band_combo_box_1.addItem(str(index))
            self.ui.band_combo_box_2.addItem(str(index))
            self.ui.band_combo_box_3.addItem(str(index))

    def update_image(self):
        if hasattr(self, 'raster_data'):
            try:
                band_1_index = int(self.ui.band_combo_box_1.currentText())
                band_2_index = int(self.ui.band_combo_box_2.currentText())
                band_3_index = int(self.ui.band_combo_box_3.currentText())
                band_indices = [band_1_index, band_2_index, band_3_index]
                image_data = self.read_bands(band_indices)
                self.display_image(image_data)
            except Exception as e:
                self.ui.info_label.setText(f'显示图像出现错误: {str(e)}')
        else:
            self.ui.info_label.setText('请先打开遥感栅格数据。')

    def read_bands(self, band_indices):
        """
        根据给定的波段索引，从栅格数据中读取对应波段的数据，并进行归一化等预处理（简单示例）
        """
        image_data = []
        for index in band_indices:
            band_data = self.raster_data.read(index).astype(np.float32)
            # 简单归一化，将数据映射到0-1范围，可根据实际优化
            band_data = (band_data - band_data.min()) / (band_data.max() - band_data.min())
            image_data.append(band_data)
        return np.dstack(image_data)

    def display_image(self, image_data):
        """
        使用matplotlib显示图像到界面的可视化展示区域
        """
        # 清除之前可能存在的可视化内容（如果有）
        if hasattr(self, 'canvas'):
            layout = self.ui.visualization_widget.layout()
            if layout is not None:
                layout.removeWidget(self.canvas)
            self.canvas.deleteLater()
            del self.canvas
            while layout.count():
                item = layout.takeAt(0)
                widget = item.widget()
                if widget is not None:
                    widget.deleteLater()
        # 创建Figure和Axes对象用于绘制图像
        fig = Figure(figsize=(6, 4))
        ax = fig.add_subplot(111)
        ax.imshow(image_data)
        ax.axis('off')  # 关闭坐标轴显示，仅显示图像
        # 创建FigureCanvas对象用于嵌入到Qt界面
        self.canvas = FigureCanvas(fig)
        # 设置可视化展示区域的布局并添加FigureCanvas
        layout = self.ui.visualization_widget.layout()
        if layout is None:
            layout = QVBoxLayout(self.ui.visualization_widget)
            layout.setContentsMargins(0, 0, 0, 0)
        else:
            # 如果布局已存在，先清空布局中的已有部件（避免重复添加）
            while layout.count():
                item = layout.takeAt(0)
                widget = item.widget()
                if widget is not None:
                    widget.deleteLater()
        layout.addWidget(self.canvas)
        visualization_widget = self.ui.visualization_widget
        visualization_widget.setLayout(layout)  # 设置布局到可视化组件
        visualization_widget.update()  # 更新可视化组件显示
        self.update()  # 更新整个窗口显示


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())